#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>

#define TCB_SIZE (64*1024) // 64KB

// 声明任务切换函数（在汇编文件中实现）
extern __attribute__((sysv_abi)) void task_switch(void* next);

// 定义64位栈帧结构，与汇编代码中的保存/恢复顺序对应
typedef struct {
    uint64_t r15;        // 保存的寄存器
    uint64_t r14;        // 保存的寄存器
    uint64_t r13;        // 保存的寄存器
    uint64_t r12;        // 保存的寄存器
    uint64_t rbx;        // 保存的寄存器
    uint64_t rbp;        // 保存的帧指针
    uint64_t ret_addr;   // 返回地址
} StackFrame;

// 任务控制块结构 - 确保总大小正好是4KB
typedef struct {
    uint64_t sp;                            // 栈指针
    uint8_t stack[TCB_SIZE - sizeof(uint64_t)]; // 栈空间
} __attribute__((aligned(TCB_SIZE))) TCB;

// 全局变量
TCB* current_task = NULL;  // 当前运行的任务
TCB* task1 = NULL;         // 任务1
TCB* task2 = NULL;         // 任务2

// 调试函数：打印任务信息
void print_task_info(TCB* task, const char* name) {
    printf("Task %s: TCB at %p, SP = %p\n", name, (void*)task, (void*)task->sp);
    
    if (task->sp != 0) {
        StackFrame* frame = (StackFrame*)(task->sp);
        printf("  Stack frame at %p:\n", (void*)frame);
        printf("  r15 = 0x%016lx\n", frame->r15);
        printf("  r14 = 0x%016lx\n", frame->r14);
        printf("  r13 = 0x%016lx\n", frame->r13);
        printf("  r12 = 0x%016lx\n", frame->r12);
        printf("  rbx = 0x%016lx\n", frame->rbx);
        printf("  rbp = 0x%016lx\n", frame->rbp);
        printf("  ret_addr = 0x%016lx\n", frame->ret_addr);
    }
}

// 初始化任务
TCB* task_create(void (*entry)(void)) {
    // 分配并初始化TCB
    TCB* task = (TCB*)aligned_alloc(4096, sizeof(TCB));
    if (!task) {
        perror("Failed to allocate TCB");
        exit(EXIT_FAILURE);
    }
    
    // 清零TCB
    for (size_t i = 0; i < sizeof(TCB); i++) {
        ((uint8_t*)task)[i] = 0;
    }
    
    // 计算栈顶位置（栈向下增长）
    uintptr_t stack_top = (uintptr_t)task->stack + sizeof(task->stack);
    
    // 确保栈对齐到16字节（x86-64 ABI要求）
    stack_top &= ~0xF;
    
    // 预先在栈上构建栈帧
    stack_top -= sizeof(StackFrame);
    StackFrame* frame = (StackFrame*)stack_top;
    
    // 初始化栈帧
    frame->r15 = 0;
    frame->r14 = 0;
    frame->r13 = 0;
    frame->r12 = 0;
    frame->rbx = 0;
    frame->rbp = 0;
    frame->ret_addr = (uint64_t)entry;  // 返回地址设为任务入口
    
    // 保存栈指针
    task->sp = stack_top;
    
    return task;
}

// 启动任务调度
void scheduler_start(TCB* first_task) {
    if (!first_task) {
        fprintf(stderr, "Cannot start scheduler with NULL task\n");
        exit(EXIT_FAILURE);
    }
    
    printf("Starting scheduler with first task at %p\n", (void*)first_task);
    print_task_info(first_task, "first");
    
    task_switch(first_task);
    
    // 不应该返回到这里
    fprintf(stderr, "Error: task_switch returned unexpectedly\n");
    exit(EXIT_FAILURE);
}

// 切换到下一个任务
void scheduler_yield(TCB* next_task) {
    if (!next_task) {
        fprintf(stderr, "Cannot yield to NULL task\n");
        return;
    }
    
    if (!current_task) {
        fprintf(stderr, "Current task is NULL, cannot yield\n");
        return;
    }
    
    printf("Yielding from task %p to task %p\n", (void*)current_task, (void*)next_task);
    print_task_info(current_task, "from");
    
    task_switch(next_task);
    
    printf("switch to task %p\n", (void*)current_task);
    print_task_info(current_task, "to");
}

void task2_func(void);
// 示例任务函数
void task1_func(void) {
    printf("Task 1 started\n");
    
    for (int i = 0; i < 5; i++) {
        printf("Task 1 running (%d)\n", i);
        //scheduler_yield(task2);  //show switch debug info
        task_switch(task2);
    }
    
    printf("Task 1 finished\n");
    exit(0);  // 结束程序
}

void task2_func(void) {
    printf("Task 2 started\n");
    
    for (int i = 0; i < 10; i++) {
        printf("Task 2 running (%d)\n", i);
        //scheduler_yield(task1);  //show switch debug info
        task_switch(task1);

    }
    
    printf("Task 2 finished\n");
}

int main() {
    printf("Cooperative scheduler demo (64-bit)\n");
    
    // 创建第一个任务
    task1 = task_create(task1_func);
    // 创建第二个任务
    task2 = task_create(task2_func);
    
    // 启动调度器
    scheduler_start(task1);
    
    // 不应该到达这里
    return 0;
}
